import numpy as np
import matplotlib.pyplot as plt

# 1.	使用K-means底层，完成以下操作
# 数据集[[1, 2], [2, 2], [6, 8], [7, 8]]
# (1)	将数据集使用合适方式进行处理（10分）
X = np.array([[1, 2], [2, 2], [6, 8], [7, 8]], dtype=np.float64)
m, n = X.shape
k = 2

# (2)	使用随机方式，初始化质心（10分）
np.random.seed(666)
rnd_idx = np.random.permutation(m)
rnd_idx = rnd_idx[:k]
C = X[rnd_idx]

# (3)	完成模型计算
# ①	设置循环5次（10分）
iters = 0
while(iters < 5):
    iters += 1

    # ②	计算数据集中每个样本点到各个质心的距离（10分）
    # ③	添加到列表中（10分）
    # ④	计算每个样本到质心距离的远近（10分）
    dis_mat = np.zeros([m, k])
    for i in range(k):
        dis_mat[:, i] = ((X - C[i])**2).sum(axis=1) ** 0.5

    # ⑤	并且将最近的质心进行记录（10分）
    class_vector = dis_mat.argmin(axis=1)

    # ⑥	使用均值更新质心（10分）
    for i in range(k):
        C[i] = X[class_vector == i].mean(axis=0)

# (4)	打印样本的归属于哪一个簇（10分）
cmap = plt.cm.get_cmap('rainbow', k)
for i in range(k):
    idx = class_vector == i
    c = cmap(i)
    plt.scatter(X[idx, 0], X[idx, 1], color=c, label=f'class#{i + 1}')

# (5)	打印质心位置（10分）
for i in range(k):
    c = cmap(i)
    plt.scatter(C[i, 0], C[i, 1], color=c, marker='x', s=100)
    plt.annotate(f'center#{i + 1}', xy=[C[i, 0], C[i, 1]])

# legend
plt.legend()

# finally show all plotting
plt.show()
